基于策略的算法（2）：TRPO, PPO

TRPO和PPO是两种基于策略的算法，它们的主要区别在于TRPO是基于约束的优化算法，而PPO是基于惩罚的优化算法。这两种算法都是在前面讲的策略梯度算法的基础上，通过一些技巧来提高算法的稳定性和收敛速度。

可以理解为REINFORCE-promax

on-policy转off-policy之重要性采样

我们前面说到Policy Gradient的on-policy方法（如REINFORCE）需依赖当前策略生成的数据更新策略，这样的算法让我们无法去完成寻找“某个区域内的优解”这样的行为（因为策略更新依赖于采样，采样依赖于策略，从而无法对策略进行多次探索之后再用采样更新，不支持跨策略重用）为了解决这个问题，我们可以引入一个onpolicy算法转化为offpolicy算法的数学工具，这个技术叫重要性采样（importance sampling）。但是在下面的PPO,TRPO之中，依然是on-policy的算法, 这里的重要性采样用于在采样之外探索更多的策略，采样重用从而提高效率，具体而言是暂时让“旧策略”去进行采样, 然后在“新策略”下计算这些数据的期望。

注意概念区分：on/offline与on/off policy

on/offline：offline是指采样和训练可以完全分开, 也就是说我们可以先采样一大堆数据, 然后再训练。而online则是指采样和训练同时进行, 也就是说我们每次采样一个数据, 就训练一次。

on/off policy：onpolicy是指我们在训练的时候, 采样的数据是由当前策略产生的。而offpolicy则是指我们在训练的时候, 采样的数据是由另一个策略产生的。off-policy和offline算法虽然看起来好像差不多, 但实际上offline要控制的分布变化和大方差比offpolicy大很多, 需要专门的设计, 例如Q-learning之中直接学maxQ, 又或者结合重要性采样, 策略约束等多种手段

重要性采样的核心思想是，我们可以在一个策略下采样数据，然后在另一个策略下计算这些数据的期望。只需要补充一个权重因子。

$E_{X\sim p1}[f(x)]\approx \frac{1}{n}\sum_{i=1}^n x_i$ 在p1下采样一组数据

$E_{X\sim p2}[f(x)]\approx \frac{1}{n}\sum_{i=1}^n \frac{p2(x)}{p1(x)}x_i$ 在p2下，只要乘以$\frac{p2(x)}{p1(x)}$就可以让两个分布的期望相等

这个方法之所以很有用在于, 我们实际上并不能表达出p1, p2(他们是$\pi(s,a)$，神经网络), 但我们可以通过采样的方法来进行这个分布的转换

但需要注意的是方差很可能会变大，所以需要一些技巧来减小方差。这就引入了我们接下来的TRPO和PPO算法。

TRPO(Trust Region Policy Optimization)

先讲什么是Trust Region。 数学定义是$N(\theta_{old})=\{\theta|\space||\theta-\theta_{old}||_2 < \delta\}$, trust region $L(\theta|\theta_{old})$在$N(\theta_{old})$内均接近目标$J(\theta)$, $L(\theta | \theta_{old})$就称为$J(\theta)$在$\theta_{old}$的trust region。

我们的目标是减少策略迭代的方差, 让策略变化不要太大，那我们就可以在当前策略的附近找到一个trust region，在这次更新之中，我们只能在这个trust region内更新策略。这样就可以保证更新的策略不会变化太大，从而减小方差。也就是全局的Find $\theta^* =argmax_{\theta}J(\theta)$ 变为了在一个小范围内Find $\theta^* =argmax_{\theta}J(\theta), \theta \in L(\theta|\theta_{old})$

结合两个思想, 我们就可以写出$J(\theta)$:

$J(\theta)= E_{S}E_A[\frac{\pi_{\theta}(a|s)}{\pi_{\theta_{old}}(a|s)}Q_{\pi}(s,a)]$ 这里用重要性采样把$a \sim \pi(\cdot|s)$消去了

而TRPO的目标是带约束的优化问题，我们的目标是找到一个$\theta$使得$J(\theta)$最大，同时满足KL散度的约束，即

Find $\theta^* =argmax_{\theta}J(\theta), \theta \in L(\theta|\theta_{old})$ and $\frac{1}{n}\sum_{i=1}^{n}D_{KL}(\pi(\cdot|s_i;\theta_{old})| \pi(\cdot|s_i;\theta)) < \delta$

这个约束问题的求解是通过共轭梯度法求Fisher矩阵之类的数学方法求解的, anyway, 这是个纯数学问题，这里不展开讲

也因为这个约束问题的求解比较麻烦, TRPO实际使用并不多

至于其他的的部分, TRPO和REINFORCE是一样的

伪代码可能如下:

1. 从经验回放池取探索性策略得到$\{s_1,a_1,r_1,s_2,a_2,r_2,...,s_T,a_T,r_T\}$一个完整的MC episode
2. 计算$u_i = \sum_{j=i}^{T}\gamma^{j-i}r_j$，这是一个MC的估计
3. Approximation: $L(\theta|\theta_{old}) = E_{S}E_A[\frac{\pi_{\theta}(a|s)}{\pi_{\theta_{old}}(a|s)}Q_{\pi}(s,a)] \approx \sum_{i=1}^{T} \frac{\pi_{\theta}(a_i|s_i)}{\pi_{\theta_{old}}(a_i|s_i)}u_i$
4. Maximization: $\theta_{new} \leftarrow argmax_{\theta}L(\theta|\theta_{old}) \space and \space ||\theta - \theta_{old}||<\delta$ 不断反向传播
5. 回到1, 清空经验回放池，重新基于新的策略采样

PPO(Proximal Policy Optimization)

主播主播, 你的TRPO计算还是太复杂了, 有没有更简单更有效的方法呢?

有的兄弟, 有的。像这样的方法，我们还有9个，其中一个就是PPO

PPO的思想很简单, 把TRPO的硬约束转换成一个软约束, 即在我们优化的$J(\theta)$里面加一项$f(\theta, \theta_{old})$作为惩罚项

同时呢，把先前的baseline也加上去，这样可以减小方差

常见的有两种加法:

1. KL散度(PPO with Adaptive KL Penalty)

$f(\theta, \theta_{old}) = \beta E_{S}D_{KL}(\pi(\cdot|s;\theta_{old})| \pi(\cdot|s;\theta))$

$L^{KLPEN(\theta)}= E_{S}E_A[\frac{\pi_{\theta}(a|s)}{\pi_{\theta_{old}}(a|s)}A_{\pi}(s,a)] - \beta E_{S}D_{KL}(\pi(\cdot|s;\theta_{old})| \pi(\cdot|s;\theta))$

$\beta$是一个超参数, 既然不好调, 那就让他自适应变化

如果KL散度>目标值, 那么$\beta \leftarrow 2\beta$, 反之$\beta \leftarrow \frac{1}{2}\beta$

2. Clip(PPO with Clip)

$f(\theta, \theta_{old}) = min(\frac{\pi_{\theta}(a|s)}{\pi_{\theta_{old}}(a|s)}, clip(\frac{\pi_{\theta}(a|s)}{\pi_{\theta_{old}}(a|s)}, 1-\epsilon, 1+\epsilon))$

$L^{CLIP(\theta)}= E_{S}E_A[min(\frac{\pi_{\theta}(a|s)}{\pi_{\theta_{old}}(a|s)}A_{\pi}(s,a), clip(\frac{\pi_{\theta}(a|s)}{\pi_{\theta_{old}}(a|s)}, 1-\epsilon, 1+\epsilon)A_{\pi}(s,a)]$

clip函数的表达式是$clip(x, a, b) = min(max(x, a), b)$, 即把x限制在[a,b]之间

而为什么要额外取一个min呢？这个min在$A_{\pi}(s,a) > 0$的时候没有影响，但在$A_{\pi}(s,a) < 0$的时候，允许使用左边的项, 更多地更新策略

• 这也是一个trick, 一个解释是, 我们希望在$A>0$的时候(策略好), 更新的不要那么大(PPO的基础), 但是在$A<0$的时候(策略不好), 我们希望更新的大一点, 允许一个大的回撤
• ref: https://stackoverflow.com/questions/46422845/what-is-the-way-to-understand-proximal-policy-optimization-algorithm-in-rl

实验表明CLIP PPO的效果更好, 所以下面给出的代码也是CLIP PPO的代码

其代码大致如下

import torch
import torch.nn as nn
import torch.optim as optim
import gym

# 定义策略网络
class PolicyNetwork(nn.Module):
    def __init__(self, state_dim, action_dim):
        super(PolicyNetwork, self).__init__()
        self.fc1 = nn.Linear(state_dim, 128)
        self.fc2 = nn.Linear(128, action_dim)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

    def act(self, state):
        logits = self.forward(state)
        probs = torch.softmax(logits, dim=-1) # 策略概率
        dist = torch.distributions.Categorical(probs)
        action = dist.sample()
        return action, dist.log_prob(action), probs


# 定义值网络
class ValueNetwork(nn.Module):
    def __init__(self, state_dim):
        super(ValueNetwork, self).__init__()
        self.fc1 = nn.Linear(state_dim, 128)
        self.fc2 = nn.Linear(128, 1)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x


# 超参数
learning_rate = 3e-4
gamma = 0.99
clip_epsilon = 0.2
num_epochs = 10
num_episodes = 1000
batch_size = 64


# 创建环境
env = gym.make("CartPole-v1")
state_dim = env.observation_space.shape[0]
action_dim = env.action_space.n

# 实例化策略网络、值网络和优化器
policy_net = PolicyNetwork(state_dim, action_dim)
value_net = ValueNetwork(state_dim)
optimizer_policy = optim.Adam(policy_net.parameters(), lr=learning_rate)
optimizer_value = optim.Adam(value_net.parameters(), lr=learning_rate)

# 训练循环
for episode in range(num_episodes):
    state = env.reset()[0]
    state = torch.tensor(state, dtype=torch.float32)
    episode_rewards = []
    log_probs = []
    old_probs = []
    states = []
    actions = []

    done = False
    while not done:
        action, log_prob, probs = policy_net.act(state.unsqueeze(0))
        next_state, reward, terminated, truncated, _ = env.step(action.item())
        done = terminated or truncated

        episode_rewards.append(reward)
        log_probs.append(log_prob)
        old_probs.append(probs)
        states.append(state)
        actions.append(action)
        state = torch.tensor(next_state, dtype=torch.float32)

    # 计算回报
    returns = []
    R = 0
    for r in episode_rewards[::-1]:
        R = r + gamma * R
        returns.insert(0, R)
    returns = torch.tensor(returns, dtype=torch.float32)

    # 计算优势函数 (这里使用简单的回报作为优势函数，可以替换为更复杂的估计方法)
    advantages = returns - value_net(torch.stack(states)).squeeze(-1)

    # PPO 损失函数
    log_probs = torch.stack(log_probs)
    old_probs = torch.stack(old_probs)
    ratio = torch.exp(log_probs - torch.log(old_probs.gather(1, actions.unsqueeze(1)).squeeze(1))) # 重要性采样比

    surrogate1 = ratio * advantages
    surrogate2 = torch.clamp(ratio, 1 - clip_epsilon, 1 + clip_epsilon) * advantages # 裁剪策略

    loss_policy = -torch.mean(torch.min(surrogate1, surrogate2)) # PPO 损失函数

    loss_value = nn.MSELoss()(value_net(torch.stack(states)).squeeze(-1), returns)

    # 更新参数
    optimizer_policy.zero_grad()
    loss_policy.backward()
    optimizer_policy.step()

    optimizer_value.zero_grad()
    loss_value.backward()
    optimizer_value.step()

    print(f"Episode: {episode+1}, Total Reward: {sum(episode_rewards)}, Policy Loss: {loss_policy.item()}, Value Loss: {loss_value.item()}")

env.close()

还可以把这玩意并行化

(the approximate param values are: K = 3-15, M = 64-4096, T (horizon) = 128-2048):